# 证型的输入和输出数据预处理
import pandas as pd
from sklearn.model_selection import train_test_split

df = pd.read_excel('./data/liverpea100.xlsx', engine='openpyxl')

# Shuffle the DataFrame
df_shuffled = df.sample(frac=1, random_state=1)  # Set random_state for reproducibility

# Split the shuffled data into training and testing datasets
train_df, test_df = train_test_split(df_shuffled, test_size=0.1, random_state=1)


def remove_spaces(s):
    return ''.join(str(s).split())


def create_txt(df, fp):
    file = open(fp, "w")
    for index, row in df.iterrows():
        use_col_ids = [1, 2, 4, 8, 9]
        inp = ''
        for i in use_col_ids:
            inp += remove_spaces(row[i])
            inp += ' '
        for i in range(13, 185):
            if pd.isna(row[i]):
                inp += '0'
            else:
                v = remove_spaces(row[i])
                try:
                    inp += str(int(float(v)))
                except:
                    inp += v
            inp += ' '
        oup = remove_spaces(row[185])

        if '脾虚湿热证' in oup:
            oup = str(0)
        elif '湿热内蕴证' in oup:
            oup = str(1)
        elif '痰瘀互结证' in oup:
            oup = str(2)
        else:
            oup = str(3)

        # Write data to the file
        file.write(inp)
        file.write("\t")
        file.write(oup)
        file.write("\n")

    # Close the file
    file.close()


create_txt(train_df, "./data/liverpea_train.txt")
create_txt(test_df, "./data/liverpea_test.txt")
